In [ ]:
import torch
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import pandas as pd
import networkx as nx
from gtda.graphs import KNeighborsGraph
import itertools
from utils import GraphRNN_dataset, GraphRNN_DataSampler
from importlib import reload
import matplotlib.pyplot as plt
print("Gathering data...")
flow_dataset = "../data/daily_county2county_2019_01_01.csv"
epi_dataset = "../data_epi/epidemiology.csv"
epi_dates = ["2020-06-09", "2020-06-10", "2020-06-11", "2020-06-12",
"2020-06-13", "2020-06-14", "2020-06-15", "2020-06-16",
"2020-06-17", "2020-06-18", "2020-06-19", "2020-06-20",
"2020-06-21", "2020-06-22", "2020-06-23", "2020-06-24",
"2020-06-25", "2020-06-26", "2020-06-27", "2020-06-28",
"2020-06-29", "2020-06-30", "2020-07-01", "2020-07-02",
"2020-07-03", "2020-07-04", "2020-07-05", "2020-07-06",
"2020-07-07", "2020-07-08", "2020-07-09", "2020-07-10",
"2020-07-11", "2020-07-12", "2020-07-13", "2020-07-14",
"2020-07-15", "2020-07-16", "2020-07-17", "2020-07-18",
"2020-07-19", "2020-07-20", "2020-07-21", "2020-07-22",
"2020-07-23", "2020-07-24", "2020-07-25", "2020-07-26",
"2020-07-27", "2020-07-28", "2020-07-29", "2020-07-30"
]
epi_dates_one_sample = ["2020-06-09", "2020-06-10", "2020-06-11", "2020-06-12",
"2020-06-13", "2020-06-14", "2020-06-15"] # for testing
epi_dates_pred_one_sample = ["2020-06-16"] # for testing
input_hor = 7
pred_hor = 1
# CREATION OF DATALOADER
from preprocessor import Preprocessor
from preprocessor import draw_network
from preprocessor import get_adj_from_plot
# INPUT TO FUNCTION: epi_dates, input_hor, pred_hor, flow_dataset, epi_dataset, SPLIT
# OUTPUT TO FUNCTION: training_data, testing_data
nan_flag_w_matrix = 0
nan_flag_train_gr_sig = 0
nan_flag_test_gr_sig = 0
def create_dataloader(epi_dates, input_hor, pred_hor, flow_dataset, epi_dataset):
training_data = []
testing_data = []
# Per input horizon and prediction horizon, we do the following:
for i in range(0, len(epi_dates) - input_hor - pred_hor):
epi_dates_one_sample = epi_dates[i:i+input_hor]
epi_dates_pred_one_sample = epi_dates[i+input_hor:i+input_hor+pred_hor]
print("Building the kronecker graph...")
# Pull in the sparsified version of the graph from Victor's preprocessor.py file
# Create a 7 day kronecker with no edge weight info or node info
preprocessor = Preprocessor(flow_dataset, epi_dataset, epi_dates_one_sample, plottable=True)
graph_df = preprocessor.combined_manual_kronecker()
# draw_network(graph_df)
print("Getting the adjacency matrix..")
w_adj_kron = get_adj_from_plot(graph_df) # RETURN THIS FOR MODEL INPUT
print("Getting the current training ex graph signal..")
tr_epi = preprocessor.set_timestep_offset_epi_dataset(from_timestep=0).get_epi_dataset()
train_graph_sig = {}
for data in tr_epi:
train_graph_sig = {**train_graph_sig, **dict(zip(data.geoid_o, data.new_confirmed))} # RETURN THIS AS GRAPH SIGNAL OF KRONECKER INPUT
print("Getting the predicted graph signal..")
pred_preproc = Preprocessor(flow_dataset, epi_dataset, epi_dates_pred_one_sample, plottable=True)
pred_epi = pred_preproc.set_timestep_offset_epi_dataset(from_timestep=0).get_epi_dataset() # returns a list of all the preprocessed epi datasets
# create a dictionary where the key is the geoid_o and the value is the new_confirmed of pred_epi -- this is the output of the model
pred_graph_sig = dict(zip(pred_epi[0].geoid_o, pred_epi[0].new_confirmed)) # RETURN THIS AS GRAPH SIGNAL OUTPUT
# if the length of pred_graph_sig is not equal to number of nodes, then raise an error
if len(pred_graph_sig) != w_adj_kron.shape[0]:
raise ValueError("The length of the predicted graph signal is not equal to the number of nodes in the graph")
print(w_adj_kron, train_graph_sig, pred_graph_sig)
training_data.append((w_adj_kron, train_graph_sig, pred_graph_sig))
# take the last 2 entries from training_data, remove them from training_data and add them to testing_date
testing_data = training_data[-2:]
training_data = training_data[:-2]
return training_data, testing_data
training_data, testing_data = create_dataloader(epi_dates, input_hor, pred_hor, flow_dataset, epi_dataset)
# store the training_data and testing_data in a pickle file
import pickle
with open("training_data.pkl", "wb") as f:
pickle.dump(training_data, f)
with open("testing_data.pkl", "wb") as f:
pickle.dump(testing_data, f)
# COMMENTS
# assuming the graph signal we are trying to predict is "new_confirmed"
# and where N is the timestep you want the offset to start. i.e., with 8,
# all nodes in the first dataset will be indexed 80..001, 800...002 etc,
# and the dataset after will be 90..001, 900...002, etc
# adjacency shape is good
# adjacency is weighted which is also good